import os
import sys
import numpy as np
import shutil
from tqdm import tqdm
from PIL import Image
import imagehash

# 从命令行参数读取源目录和like文件夹
if len(sys.argv) < 2:
    print("使用方法: python script.py <source_dir> <like_dir>")
    sys.exit(1)

source_dir = sys.argv[1]  # 输入文件夹
like_dir = "like2"   # like文件夹
out_dir = os.path.join(source_dir, "__out__")  # 输出文件夹

# 确保目标文件夹存在
os.makedirs(out_dir, exist_ok=True)

# 计算图片的感知哈希值 (pHash)
def calculate_image_hash(image):
    return imagehash.phash(image)

# 计算两个图片哈希值之间的相似度，返回汉明距离
def calculate_similarity(hash1, hash2):
    return 1 - (hash1 - hash2) / len(hash1.hash) ** 2  # 相似度 = 1 - 汉明距离

# 裁剪遮挡区域
def crop_face(image, box):
    (x, y, x1, y1) = box.astype("int")
    return image[y:y1, x:x1]

# 用于存储like文件夹中的图片哈希值
like_image_hashes = []

# 从like文件夹读取所有图片并计算哈希值
for root, dirs, files in os.walk(like_dir):
    for filename in files:
        if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
            img_path = os.path.join(root, filename)
            img = Image.open(img_path)
            image_hash = calculate_image_hash(img)
            like_image_hashes.append((image_hash, img_path))

# 处理进度条
total_files = sum([len(files) for r, d, files in os.walk(source_dir)])

# 统计信息
similar_images_count = 0
ok_file_path = os.path.join(out_dir, "ok.txt")

# 处理进度条
with tqdm(total=total_files, desc="处理进度", unit="file", ncols=100) as pbar:
    for root, dirs, files in os.walk(source_dir):
        for filename in files:
            if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
                img_path = os.path.join(root, filename)

                try:
                    # 读取图像
                    img = Image.open(img_path)

                    # 这里需要定义你的遮挡框，假设用 (x, y, x1, y1) 作为示例
                    face_box = np.array([50, 50, 150, 150])  # 替换为你的实际遮挡框
                    
                    # 裁剪掉遮挡部分
                    cropped_img = crop_face(np.array(img), face_box)

                    # 将裁剪后的 NumPy 数组转换回 PIL 图像
                    cropped_img = Image.fromarray(cropped_img)

                    # 计算裁剪后的图像哈希值
                    image_hash = calculate_image_hash(cropped_img)
                    similar_image_found = False

                    # 检查图片是否与like文件夹中的图片相似
                    for like_hash, like_image_path in like_image_hashes:
                        similarity = calculate_similarity(image_hash, like_hash)
                        if similarity > 0.7:  # 设定相似度阈值为70%
                            shutil.copy(img_path, os.path.join(out_dir, filename))
                            similar_images_count += 1  # 增加相似图片计数

                            # 立即写入 ok.txt 文件
                            with open(ok_file_path, 'a', encoding='utf-8') as f:
                                f.write(f"{img_path} 与 {like_image_path} 相似\n")

                            tqdm.write(f"找到相似: {img_path} 和 {like_image_path} (相似度: {similarity:.2f})")
                            similar_image_found = True
                            break  # 找到相似图片后跳出循环

                except OSError as e:
                    tqdm.write(f"跳过损坏的图像文件: {img_path}，错误: {e}")

                pbar.update(1)

print(f"处理完成！找到相似的图片总数: {similar_images_count}。\n所有相似图片已保存到目标目录: {out_dir}")
